import os.path
from docx.document import Document as _Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import _Cell, Table
from docx.text.paragraph import Paragraph

import os
from config import SUPPORTED_IMAGE_FORMATS
from pasers.paser_utils import combine_images, re_filter


# 从block中提取图片文件
def process_images(figure_block, doc, doc_id, block_id, img_save_path, content, image_info):
    """
    处理图像并提取信息。

    Args:
        figure_block: 文档中的图像块
        doc: 当前文档对象
        doc_id: 文档ID
        block_id: 块ID
        img_save_path: 图像保存路径
        content: 内容列表
        image_info: 图像信息字典
    """
    img_elements = figure_block._element.xpath('.//pic:pic')
    if not img_elements:
        return

    if len(img_elements) == 1:
        # 单个图像处理
        process_single_image(img_elements[0], doc, doc_id, block_id, img_save_path, content, image_info)
    else:
        # 多个图像处理
        process_multiple_images(img_elements, doc, doc_id, block_id, img_save_path, content, image_info)


def process_single_image(img_element, doc, doc_id, block_id, img_save_path, content, image_info):
    """处理单个图像的存储与信息提取"""
    embed = img_element.xpath('.//a:blip/@r:embed')[0]
    related_part = doc.part.related_parts[embed]
    suffix = related_part.filename.split('.')[-1].lower()

    if suffix in SUPPORTED_IMAGE_FORMATS:
        image_name = f'{doc_id}_{block_id}_0.{suffix}'
        write_image_to_file(related_part.blob, os.path.join(img_save_path, image_name))
        image_key = f'image_{doc_id}_{block_id}_0'
        image_info[image_key] = {'path': image_name, 'content': ''}  # Placeholder for actual text extraction
        content.append(f'<image:{image_key}>')
    else:
        print(f'不能提取{suffix}图片')


def process_multiple_images(img_elements, doc, doc_id, block_id, img_save_path, content, image_info):
    """处理多个图像的存储与信息提取"""
    image_blobs = []

    for img_element in img_elements:
        embed = img_element.xpath('.//a:blip/@r:embed')[0]
        related_part = doc.part.related_parts[embed]
        suffix = related_part.filename.split('.')[-1].lower()

        if suffix in SUPPORTED_IMAGE_FORMATS:
            image_blobs.append((related_part.blob, suffix))
        else:
            print(f'不能提取{suffix}图片')

    if image_blobs:
        image_name = f'{doc_id}_{block_id}_1.{image_blobs[0][1]}'
        image_key = f'image_{doc_id}_{block_id}_1'
        combine_images(image_blobs, saved_path=os.path.join(img_save_path, image_name))
        image_info[image_key] = {'path': image_name, 'content': ''}  # Placeholder for actual text extraction
        content.append(f'<image:{image_key}>')


def write_image_to_file(image_blob, path):
    """将图像数据写入文件"""
    with open(path, 'wb') as f:
        f.write(image_blob)


def get_table_data(table_block, max_head_index=1):
    """
    table: 表格块
    max_len: 将表格切分为子表格，每个子表格的字符串长度不超过 max_len
    max_head_index: 多行表头的最大行索引，索引从0开始计数
    return: markdown字符串形式输出子表列表
    """
    row_data = []  # 表数据，其中每一行的内容总长度不超过的最大长度
    table_head = []  # 表头数据，存放表头
    head_index = -1  # 表头是否建立完成
    head_cell_num = len(table_block.rows[0].cells)  # 表头元素数量
    for i, row in enumerate(table_block.rows):
        column_types = set()
        row_text = []
        for cell_id, cell in enumerate(row.cells):
            cell_text = cell.text.replace('\n', '').strip()
            if len(cell_text) == 0:
                row_text.append('')
            else:
                row_text.append(cell_text.strip())
                column_types.add(cell_text.strip())
        if len(''.join(row_text).strip()) != 0:
            row_text = get_row_text(head_cell_num, row_text)
            cur_row_md = to_md_string(row_text)
            if head_index == -1 and len(column_types) < head_cell_num and len(table_head) < max_head_index:
                # 表头判定条件，当表头数据具有唯一性时被确认为表头
                table_head.append(cur_row_md)
            elif head_index == -1 and (len(column_types) >= head_cell_num or len(table_head) >= max_head_index):
                # 表头结束判定
                table_head.append(cur_row_md)
                table_head.append('|' + '-|' * head_cell_num)
                head_index = i
            # 表数据内容写入
            else:
                if cur_row_md.split('|')[1:-1]:
                    row_data.append(cur_row_md)
    # table_head, row_data = remove_empty_col(table_head, row_data)
    if table_head:
        return {'row_data': row_data, 'table_head': table_head}
    else:  # 表头为空说明是个空表，这种情况通常出现在pdf解析中
        return None


def get_row_text(head_cell_num, row_text):
    """
    使 row_text 中的元素数量和 head_cell_num 中数量相等
    :param head_cell_num:
    :param row_text:
    :return: row_text
    """
    if head_cell_num == len(row_text):
        return row_text
    elif head_cell_num < len(row_text):
        row_text[head_cell_num - 1] = ''.join(row_text[head_cell_num - 1:]).strip()
        return row_text[:head_cell_num]
    else:
        row_text.extend([''] * (head_cell_num - len(row_text)))
        return row_text


def to_md_string(a_ls):
    """
    将 a_ls 列表转换为 markdown 字符串
    :param a_ls:
    :return:
    """
    b = ['|']
    for i in a_ls:
        b.append(str(i).replace('|', '').strip())
        b.append('|')
    return ''.join(b)


# 遍历docx文件中所有的子元素
def iter_block_items(parent):
    if isinstance(parent, _Document):
        parent_elm = parent.element.body
    elif isinstance(parent, _Cell):
        parent_elm = parent._tc
    else:
        raise ValueError("something's not right")

    for child in parent_elm.iterchildren():
        if isinstance(child, CT_P):
            yield Paragraph(child, parent)
        elif isinstance(child, CT_Tbl):
            yield Table(child, parent)


# 文件无格式，直接输出原文
def get_content(doc_id, doc, file_type, img_save_path):
    """
    :param doc: Document对象
    :param doc_id: 文件类型标识符
    :param img_save_path: 图片保存路径
    :param file_type: 文件类型
    :return: 表图文分割列表；图文件所在路径
    """
    table_info = {}
    image_info = {}
    # exclude_image_texts = {'ET)', ''}  # 如果ocr内容属于这个集合中，就丢弃这个图片
    content_ls = []
    for block_id, block in enumerate(iter_block_items(doc)):
        if isinstance(block, Paragraph):
            if len(block.text) > 0:
                content_ls.append(block.text)
            # 图片处理，若有图片则将图片保存到img_path
            process_images(block, doc, doc_id, block_id, img_save_path, content_ls, image_info)
        elif isinstance(block, Table):
            table_data = get_table_data(block)
            if table_data is not None:
                if file_type == 'pdf':
                    row_num = len(table_data['row_data'])
                    column_num = len(table_data['table_head'][-1].split('|')) - 2
                    if row_num < 3 and column_num < 3:  # pdf表格少于2行（不包括头）2列的全部当做纯文本处理
                        content = table_data['table_head'][0].replace('|', ' ') + '\n' + "\n".join(
                            table_data['row_data']).replace('|', ' ').strip()
                        content_ls.append(content)
                else:
                    if not table_data['row_data']:
                        content_ls.append(table_data['table_head'][0].replace('|', ' '))
                    else:
                        content_ls.append(f'<table:table_{block_id}>')
                        table_info[f'table_{block_id}'] = table_data
    return re_filter('\n'.join(content_ls)), image_info, table_info


def remove_empty_col(table_head, row_data):
    # 寻找需要被删除的列与行
    non_empty_cols = set()
    non_empty_rows = set()
    for i, row in enumerate(row_data):
        cells = row.split('|')
        for j in range(len(cells)):
            if cells[j] != '':
                non_empty_cols.add(j)
                non_empty_rows.add(i)

    # 删除不需要的列
    def form_new_table(rows):
        new_data = []
        for i, row in enumerate(rows):
            if i not in non_empty_rows:
                continue
            cells = row.split('|')
            new_row = '|'.join([cell for j, cell in enumerate(cells) if j in non_empty_cols])
            new_data.append('|' + new_row + '|')
        return new_data

    new_table_head = form_new_table(table_head)
    new_row_data = form_new_table(row_data)
    return new_table_head, new_row_data
